#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence 
@file: queue_mq_thread.py
@time: 2016/12/8 13:45
"""
import json
import threading
import time
from Queue import Queue

from beanstalkc import SocketError, CommandFailed

from common.pybeanstalk import PyBeanstalk


class BeanstalkHandler(threading.Thread):
    PAUSE_BLOCK_TIME = 10

    PAUSE_NORMAL_TIME = 0.5

    # 最大待发送队列大小
    MAX_QUEUE_SIZE = 50

    def __init__(self, beanstalk_conf, max_block=20, log=None):
        threading.Thread.__init__(self)
        self.daemon = True

        self.log = log

        # 输送队列
        self.queue = Queue()

        # 最大阻塞数目
        self.__max_block = max_block

        # 获取消息队列配置
        self.server_conf = beanstalk_conf

        self.__beanstalk = PyBeanstalk(self.server_conf['host'], self.server_conf['port'])

    # 判断是否需要暂停
    def __wait_beanstalk_queue(self, tube, block):
        times = 0
        while True:
            sleep_time = self.PAUSE_NORMAL_TIME
            try:
                count = self.__beanstalk.get_tube_count(tube)
            except CommandFailed, (_, status, results):
                if status == 'NOT_FOUND':
                    self.log.error('获取消息队列不存在: tube = {}'.format(tube))
                    count = 0
                else:
                    self.log.error('获取消息队列异常: tube = {} status = {}'.format(tube, status))
                    sleep_time = self.PAUSE_BLOCK_TIME
                    count = self.__max_block
            except Exception as e:
                sleep_time = self.PAUSE_BLOCK_TIME
                count = self.__max_block
                self.log.error("获取消息队列大小失败: ")
                self.log.exception(e)

            if count < self.__max_block:
                return True

            # 如果不阻塞 则直接返回当前数据丢失
            if not block:
                return False

            times += 1
            if times % 20 == 0:
                self.log.info("当前消息队列阻塞需要休眠: tube = {} sleep = {}".format(tube, sleep_time))

            # 开始休眠
            time.sleep(sleep_time)

    # 判断内存队列是否有阻塞
    def __wait_mem_queue(self, block):

        count = 0
        while True:
            count += 1
            q_size = self.queue.qsize()
            if q_size >= self.MAX_QUEUE_SIZE:
                # 如果不阻塞等待 则直接返回错误，队列已满 不再添加数据
                if not block:
                    return False

                if count % 100 == 0:
                    self.log.info("当前线程队列阻塞, 需要休眠: queue size = {}".format(q_size))

                time.sleep(0.2)
                continue
            break
        return True

    def close(self):
        self.queue.put_nowait('beanstalk close')

    def push_msg(self, tube, item_dict, block=True):

        # 判断是否需要等待
        if not self.__wait_mem_queue(block):
            return False

        msg = json.dumps(item_dict)
        self.queue.put_nowait((tube, msg, block))
        return True

    def run(self):
        self.log.info('开始运行消息队列...')
        while True:
            try:

                body = self.queue.get()

                # 如果是字符串则直接退出，视为退出消息
                if isinstance(body, basestring):
                    self.log.info("消息队列线程收到字符串数据, 退出线程: body = {}".format(body))
                    break

                if not isinstance(body, tuple):
                    self.log.error("消息类型错误: body type = {} msg = {}".format(type(body), body))
                    time.sleep(5)
                    continue

                tube, msg, block = body
                while True:
                    try:
                        # 发送前先判断是否需要休眠 如果阻塞则直接返回
                        if not self.__wait_beanstalk_queue(tube, block):
                            break

                        result = self.__beanstalk.put(tube, msg)
                        if result is None:
                            self.log.error("数据量太大: msg len = {}".format(len(msg)))

                        break
                    except SocketError as e:
                        time.sleep(10)
                        self.__beanstalk.reconnect()
                        self.log.warn("reconnect beanstalk...")
                        self.log.exception(e)
                    except Exception as e:
                        self.log.error('捕获异常休眠:')
                        self.log.exception(e)
                        time.sleep(10)
            except Exception as e:
                self.log.info('当前队列大小: size = {size}'.format(size=self.queue.qsize()))
                self.log.exception(e)
                time.sleep(5)
        self.log.info("消息队列正常退出...")
